from openai import OpenAI
import base64
import os
import re
import time
from com.watermelon.rpa.functions import kimi_config as kimi, kimi_config

'''
client = OpenAI(
    api_key= kimi.API_KEY,
    base_url="https://api.moonshot.cn/v1",
)

completion = client.chat.completions.create(
    model="moonshot-v1-8k",
    messages=[
        {"role": "system",
         "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力等问题的回答。Moonshot AI 为专有名词，不可翻译成其他语言。"},
        {"role": "user", "content": "你好，我叫李雷，1+1等于多少？"}
    ],
    temperature=0.3,
)

print(completion.choices[0].message.content)
'''


def kimiFunction(apiKey,pngFilePath,text):
    client = OpenAI(
        api_key=apiKey,
        base_url="https://api.moonshot.cn/v1",
    )

    # 对图片进行base64编码
    with open(pngFilePath, 'rb') as f:
        img_base = base64.b64encode(f.read()).decode('utf-8')

    response = client.chat.completions.create(
        model="moonshot-v1-8k-vision-preview",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{img_base}"
                        }
                    },
                    {
                        "type": "text",
                        "text": text
                    }
                ]
            }
        ]
    )
    return response

#获取pdf文件，并返回pdf文件路径+fileName
def obtain_pdf_file(filePath):
    listArr = []
    allFile = os.listdir(filePath)
    for file in allFile:
        if os.path.isfile(os.path.join(filePath,file)) and os.path.splitext(file)[1].lower() == ".png":
            listArr.append(filePath+"\\"+file)
    return listArr

#识别图片文字，并用列表打包返回
def identify(png):
    arrTest = []
    for i in png:
        time.sleep(1)
        result = kimiFunction(kimi_config.API_KEY, i, "识别图片中的文字，返回结果不需要多余的信息，以逗号分割")
        content = result.choices[0].message.content
        arr = content.split("，")
        arrTest.append(arr)
    return arrTest


filePath = r"C:\Users\Big Watermelon\Desktop\shuiguotupian\test"

png = obtain_pdf_file(filePath)

print(identify(png))